

import os
import xmltodict
import json

def xml_to_json(xml_dir):
    json_path = os.path.dirname(xml_dir)+os.sep+"via_region_data_ori.json" 

    if os.path.isfile(xml_dir):
        xml_dirs = [xml_dir]
    else:
        xml_dirs = [xml_dir+os.sep+f for f in  os.listdir(xml_dir)]

    json_dict = dict()                              
    for xml_path in xml_dirs:
        
        with open(os.path.join(xml_dir,xml_path),'r',encoding='utf-8') as x:    
            x_str = x.read()
            xml_dic = xmltodict.parse(x_str,encoding='utf-8')    
        xml_dic = xml_dic["annotation"]
        

        one_label = dict()
        filename = xml_dic["filename"]
        if not filename.endswith(".jpg"):
            filename += ".jpg"
        one_label["filename"] =  filename

        if xml_dic.get("object", None) is None:
            print("error: ", filename)
            continue

        xml_dics = []
        if type(xml_dic["object"]) == list:
            xml_dics = xml_dic["object"]
        else:
            xml_dics.append(xml_dic["object"])

        regions = []
        for region in xml_dics:
            # print(region)
            x = int(region["bndbox"]["xmin"])
            y = int(region["bndbox"]["ymin"])
            width = int(region["bndbox"]["xmax"] )- x
            height = int(region["bndbox"]["ymax"] )- y
            rect = dict()
            rect["shape_attributes"] = dict()
            rect["shape_attributes"]["name"] = "rect"

            rect["shape_attributes"]["x"] = x
            rect["shape_attributes"]["y"] = y 
            rect["shape_attributes"]["width"] = width 
            rect["shape_attributes"]["height"] = height 

            rect["region_attributes"] = dict()
            rect["region_attributes"]["label"] = "rect" 
            regions.append(rect)           
        one_label["regions"] = regions

        json_dict[filename] = one_label
    with open(json_path, 'w', encoding="utf-8") as wf:    
        wf.write(json.dumps(json_dict) )

if __name__ == "__main__":
    xml_dir = r'C:\Users\swls\Documents\WeChat Files\A360697611\FileStorage\File\2020-01\0106'      #xml文件目录
    # xml_dir = r'C:\Users\swls\Documents\WeChat Files\A360697611\FileStorage\File\2020-01\2020.0106\11_033e6f3ab0e1960710a429be1c8bc5d8_384.xml'      #xml文件目录

    xml_to_json(xml_dir)